import os
import DEAD.AutoDecoder.Objective.Score as Score
from Utility.Xml import *

#数据字典的路径，用户可修改
dictionary_path = "DEAD/AutoDecoder/Dictionaries/dictionary_G47800.xml"

# 获取根元素
root = get_root_from_xml(dictionary_path)

# 选择模式
mode = get_int_from_xml(root, "in.selection.mode")
# 选用相应的损失函数
if mode <= 2:
    from DEAD.AutoDecoder.Objective.Loss import loss_sdf as loss_function
else:
    from DEAD.AutoDecoder.Objective.Loss import loss_eikonal as loss_function
# 是否新建loss曲线文件
create_new_curve={
    1: True,
    2: False,
    3: True,
    4: False
}

training_groups = {
    1: "training_SDF",
    2: "training_SDF_further",
    3: "training_eikonal",
    4: "training_eikonal_further"
}
training_group = training_groups[mode]


# 数据集相关参数
trainingset_path = get_str_from_xml(root, "in.dataset.trainingset_path")
testingset_path = get_str_from_xml(root, "in.dataset.testingset_path")
data_num = get_int_from_xml(root, "in.dataset.data_num")
random_point_num=get_int_from_xml(root, "in.dataset.random_point_num")

# 神经网络相关参数
hidden_size = get_int_from_xml(root, "in.network.hidden_size")
depth = get_int_from_xml(root, "in.network.depth")
latent_size = get_int_from_xml(root, "in.network.latent_size")

# 训练相关参数
model_source_path = get_str_from_xml(root, f"in.{training_group}.model_source_path")
model_save_path = get_str_from_xml(root, f"in.{training_group}.model_save_path")
os.makedirs(model_save_path, exist_ok=True)
epoch_num = get_int_from_xml(root, f"in.{training_group}.epoch_num")
batch_size = get_int_from_xml(root, f"in.{training_group}.batch_size")
learning_rate = get_double_from_xml(root, f"in.{training_group}.learning_rate")
validation_batch_size=get_int_from_xml(root, f"in.{training_group}.validation_batch_size")
validation_interval=get_int_from_xml(root, f"in.{training_group}.validation_interval")

# 评估相关参数
model_load_path = get_str_from_xml(root, "in.evaluation.model_load_path")
decode_result_save_path=get_str_from_xml(root,"in.evaluation.decode_result_save_path")

# PSFD相关参数
PSFD_save_path=get_str_from_xml(root,"in.PSFD.PSFD_save_path")
R_min=get_double_from_xml(root,"in.PSFD.R_min")
R_max=get_double_from_xml(root,"in.PSFD.R_max")
J_min=get_double_from_xml(root,"in.PSFD.J_min")
J_max=get_double_from_xml(root,"in.PSFD.J_max")
cstar_min=get_double_from_xml(root,"in.PSFD.cstar_min")
cstar_max=get_double_from_xml(root,"in.PSFD.cstar_max")
rho_min=get_double_from_xml(root,"in.PSFD.rho_min")
rho_max=get_double_from_xml(root,"in.PSFD.rho_max")
r_ref_min=get_double_from_xml(root,"in.PSFD.r_ref_min")
r_ref_max=get_double_from_xml(root,"in.PSFD.r_ref_max")
n_min=get_double_from_xml(root,"in.PSFD.n_min")
n_max=get_double_from_xml(root,"in.PSFD.n_max")

training_infos = {
    1: f"\n\033[92m模式:1 训练SDF\n\n以SDF为损失函数从头开始训练模型\n会清除之前训练的模型(\033[91m{model_save_path}\033[92m)\n在继续前务必备份\033[0m\n",
    2: f"\n\033[92m模式:2 进一步训练SDF\n\n在SDF训练结果基础上继续以SDF为损失函数训练模型\n保存路径为\033[91m{model_save_path}\033[92m\n留意是否会覆盖之前训练的模型(\033[91m{model_source_path}\033[92m)\n在继续前建议备份\033[0m\n",
    3: f"\n\033[92m模式:3 训练Eikonal\n\n在SDF训练结果基础上以Eikonal方程残差为损失函数训练模型\n保存路径为\033[91m{model_save_path}\033[92m\n留意是否会覆盖之前训练的模型(\033[91m{model_source_path}\033[92m)\n在继续前建议备份\033[0m\n",
    4: f"\n\033[92m模式:4 进一步训练Eikonal\n\n在Eikonal训练结果基础上继续以Eikonal方程残差为损失函数训练模型\n保存路径为\033[91m{model_save_path}\033[92m\n留意是否会覆盖之前训练的模型(\033[91m{model_source_path}\033[92m)\n在继续前建议备份\033[0m\n"
}
training_info = training_infos[mode]
